from os.path import join, abspath
from backend.experiment.framework.shots_prompt.task import \
    PerplexityPromptShotsClassificationTaskMaker
from backend.experiment.framework.task import TaskMaker
from backend.experiment.tokenizer import tokenizer
from backend.experiment.framework.shots_prompt.implement import \
    CommonLabelBalanceStaticShotsGenerator
from backend.experiment.model import context_limit
from backend.experiment.clue.afqmc.scheme0.implement import AFQMCImplement as \
    S0Impl
from typing import List


class AFQMCImplement(S0Impl):
    @classmethod
    def label_map(cls) -> List[str]:
        return ['不同', '相同']

    @classmethod
    def _make_prompt(cls, record, mask_target: bool) -> str:
        s1: str = record['sentence1']
        s2: str = record['sentence2']
        label: str = cls.__mask_token if mask_target else \
            cls.label_map()[int(record['label'])]
        return f'下面两个句子语义{label}：{s1}。{s2}'

    def get_task_maker(self) -> TaskMaker:
        return PerplexityPromptShotsClassificationTaskMaker(
            token_limit=context_limit,
            tokenizer=tokenizer,
            shots_generator=CommonLabelBalanceStaticShotsGenerator(
                labeled_example_tokens=self.classify_example_tokens(
                    self.get_shuffle_example_tokens()
                )
            ),
            label_map=self.label_map()
        )

    @classmethod
    def cache_dir(cls) -> str:
        return join(cls.work_dir(), 'cache')

    @classmethod
    def work_dir(cls) -> str:
        return abspath(join(__file__, '..'))
